import logging  # 用于日志记录
import time  # 用于时间相关操作
from ruamel import yaml  # 用于YAML配置文件的读写（支持保留格式）


class AverageMeter(object):
    """
    平均值计算器类：用于跟踪和计算一组数值的当前值、总和、计数及平均值
    常用于训练过程中统计损失、准确率等指标的实时变化
    """

    def __init__(self):
        self.reset()  # 初始化时重置所有统计值

    def reset(self):
        """重置所有统计变量"""
        self.val = 0  # 当前值（最新加入的数值）
        self.avg = 0  # 平均值（sum / count）
        self.sum = 0  # 总和（所有数值的累计和）
        self.count = 0  # 计数（数值的总数量）

    def update(self, val, n=1):
        """
        更新统计值
        Args:
            val: 新加入的数值（如单次batch的损失）
            n: 该数值对应的样本数量（默认1，用于加权计算）
        """
        self.val = val  # 更新当前值
        self.sum += val * n  # 累加加权总和
        self.count += n  # 累加样本数量
        self.avg = self.sum / self.count  # 重新计算平均值


def Prepare_logger(args, eval=False):
    """
    日志记录器配置函数：创建并配置日志记录器，同时输出到控制台和文件
    Args:
        args: 包含配置参数的对象（需包含snapshot_pref日志保存路径）
        eval: 是否为评估模式（用于区分训练/评估日志文件，默认False）
    Returns:
        logger: 配置好的日志记录器对象
    """
    # 获取日志记录器实例（指定名称为当前模块）
    logger = logging.getLogger(__name__)
    logger.propagate = False  # 禁止日志向上传播（避免重复输出）
    logger.setLevel(logging.INFO)  # 设置日志级别为INFO（只记录INFO及以上级别）

    # 配置控制台输出处理器
    handler = logging.StreamHandler()  # 控制台输出流
    # 定义日志格式：时间 + 级别 + 消息内容
    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
    handler.setFormatter(formatter)  # 应用格式
    handler.setLevel(0)  # 处理器级别（0表示所有级别都处理）
    logger.addHandler(handler)  # 添加到日志器

    # 生成日志文件名（包含当前时间戳，区分训练和评估日志）
    date = time.strftime('%Y%m%d%H%M', time.localtime(time.time()))  # 格式：年月日时分
    # 训练日志路径：snapshot_pref + 时间.log；评估日志路径：snapshot_pref/时间-Eval.log
    logfile = args.snapshot_pref + date + '.log' if not eval else args.snapshot_pref + f'/{date}-Eval.log'

    # 配置文件输出处理器
    file_handler = logging.FileHandler(logfile, mode='w')  # 文件输出流（覆盖写入模式）
    file_handler.setLevel(logging.INFO)  # 文件日志级别为INFO
    file_handler.setFormatter(formatter)  # 应用相同格式
    logger.addHandler(file_handler)  # 添加到日志器

    return logger


def get_configs(dataset):
    """
    获取数据集配置函数：从YAML文件中读取指定数据集的配置信息
    Args:
        dataset: 数据集名称（用于索引配置）
    Returns:
        该数据集对应的配置字典
    """
    # 读取配置文件（./configs/dataset_cfg.yaml），返回所有数据集的配置
    data = yaml.load(open('./configs/dataset_cfg.yaml'))
    return data[dataset]  # 返回指定数据集的配置


def get_and_save_args(parser):
    """
    参数处理与保存函数：解析命令行参数，与默认配置合并，并保存最终配置到文件
    Args:
        parser: 命令行参数解析器（argparse.ArgumentParser对象）
    Returns:
        default_config: 合并后的最终配置字典
    """
    args = parser.parse_args()  # 解析命令行参数

    # 读取默认配置文件（./configs/default_config.yaml），使用RoundTripLoader保留格式
    default_config = yaml.load(open('./configs/default_config.yaml', 'r'), Loader=yaml.RoundTripLoader)
    current_config = vars(args)  # 将命令行参数转换为字典

    # 合并命令行参数与默认配置：命令行参数覆盖默认配置
    for k, v in current_config.items():
        if k in default_config:  # 只处理默认配置中已有的键
            # 若命令行参数非None且与默认值不同，则更新默认配置
            if (v != default_config[k]) and (v is not None):
                print(f"Updating:  {k}: {default_config[k]} (default) ----> {v}")
                default_config[k] = v

    # 将合并后的配置保存到当前目录（./current_configs.yaml），保留格式
    yaml.dump(default_config, open('./current_configs.yaml', 'w'), indent=4, Dumper=yaml.RoundTripDumper)
    return default_config